import torchvision.transforms as T
from torchvision.datasets import CocoDetection
from torch.utils.data import DataLoader
import torchvision.models.detection as models
from torchmetrics.detection.mean_ap import MeanAveragePrecision
from tqdm import tqdm
import numpy as np
import torch
import json


def collate_fn(batch):
    """
    解决 RuntimeError: each element in list of batch should be of equal size
    """
    data, targets = zip(*batch)  # 解包批次中的数据和标签
    new_targets = []
    for objs in targets:
        boxes = torch.tensor([obj['bbox'] for obj in objs], dtype=torch.float32)  # 获取边界框
        labels = torch.tensor([obj['category_id'] for obj in objs], dtype=torch.long)  # 获取类别标签
        new_targets.append({'boxes': boxes, 'labels': labels})
    return data, new_targets

    # modified_targets = [t + 1 for t in targets]  # 例如，将所有标签加1
    # new_targets = []

    # for objs in targets:
    #     # print(objs)
    #     labels = []
    #     boxes = []
    #     if objs is not None and len(objs) > 0:
    #         for obj in objs:
    #             category_id = obj['category_id']
    #             bbox = obj['bbox']
    #             labels.append(category_id)
    #             boxes.append(bbox)
    #     objs = {
    #         'labels': torch.tensor(labels),
    #         'boxes': torch.tensor(boxes)
    #     }
    #     # print(objs)
    #     new_targets.append(objs)

    # return data, new_targets
dataset = []
def get_dataloader():
    """
    获取数据加载器
    """
    # 转换操作，例如调整大小和归一化
    transform = T.Compose([
        T.Resize((300, 300)),  # SSD通常使用固定大小的输入图像
        T.ToTensor()
    ])
    
    dataset_path = '/workspace/pytorch_performance_metrics/data/coco128'
    global dataset
    # 加载数据集
    dataset = CocoDetection(root=f'{dataset_path}/images/train2017', annFile=f'{dataset_path}/train2017_coco_format.json', transform=transform)
    dataloader = DataLoader(dataset, batch_size=4, shuffle=False, num_workers=4, collate_fn=collate_fn)

    return dataloader

def postprocess(predictions, score_threshold=0.5):
    # 对每个图像的预测结果进行后处理
    for i in range(len(predictions)):
        # 过滤置信度低于阈值的检测结果
        scores = predictions[i]['scores']
        boxes = predictions[i]['boxes']
        labels = predictions[i]['labels']
        
        # 保留置信度高于阈值的检测结果
        keep = torch.where(scores > score_threshold)[0]
        boxes = boxes[keep]
        scores = scores[keep]
        labels = labels[keep]
        
        # # 应用非极大值抑制
        # keep = torchvision.ops.nms(boxes, scores, iou_threshold)
        # boxes = boxes[keep]
        # scores = scores[keep]
        # labels = labels[keep]
        
        # 更新预测结果
        predictions[i]['boxes'] = boxes
        predictions[i]['scores'] = scores
        predictions[i]['labels'] = labels
    return predictions


# 绘制检测结果
def draw_detections(image, predictions, class_labels):
    # image = image.copy()
    boxes = predictions['boxes'].numpy()
    labels = predictions['labels'].numpy()
    scores = predictions['scores'].numpy()
    
    for box, label, score in zip(boxes, labels, scores):
        x1, y1, x2, y2 = map(int, box)
        class_name = class_labels[label]
        
        # 绘制边界框
        # cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        # 添加类别标签和置信度
        text = f"{class_name}: {score:.2f}"
        # cv2.putText(image, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    return image

# COCO 数据集的类别标签
COCO_CLASS_LABELS = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
    'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
    'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]

def measure_performance(dataloader, model, model_type):
    """
    评估性能指标
    """
    # 初始化评估指标
    iou_thresholds = np.linspace(0.5, 0.95, 10).tolist()  # IoU阈值范围从0.5到0.95，步长为0.05
    mean_ap = MeanAveragePrecision(iou_thresholds=iou_thresholds)
    # precision_recall = DetectionPrecisionRecall(iou_thresholds=iou_thresholds)

    all_targets = []
    all_preds = []
    with torch.no_grad():
        for images, targets in tqdm(dataloader):
            # print(targets)
            # images = list(image.to("cuda") for image in images)
            # targets = [{k: v.to("cuda") for k, v in t.items()} for t in targets]
            outputs = model(images)
<<<<<<< HEAD
            # outputs = postprocess(outputs)  

            # 绘制检测结果
            # result_image = draw_detections(images[0], outputs[0], COCO_CLASS_LABELS)

            # # 显示结果
            # cv2.imshow('Detection Result', result_image)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

            # # 保存结果
            # cv2.imwrite('./detection_result.jpg', result_image)
=======
            # print(outputs)
            # print(targets)
>>>>>>> cc3e4d15239c6bef0e6c3f157ab69d4ba6637c29
            all_preds.extend(outputs)
            new_targets = []
            for objs in targets:
                # print(objs)
                labels = []
                boxes = []
                if objs is not None and len(objs) > 0:
                    for obj in objs:
                        category_id = obj['category_id']
                        bbox = obj['bbox']
                        labels.append(category_id)
                        boxes.append(bbox)
                new_targets.append({
                    'labels': torch.tensor(labels),
                    'boxes': torch.tensor(boxes)
                })
            print('new_targets')    
            print(new_targets)
            all_targets.extend(new_targets)
            # 可以使用以下代码片段更新指标（如果需要）: 注意这种方式通常不直接用于更新指标，而是在所有数据后汇总计算。
            # mean_ap.update(preds=outputs, target=targets)
            # precision_recall.update(preds=outputs, target=targets)
    
    mean_ap_final = mean_ap(all_preds, all_targets)  # 注意这里的用法可能需要调整，具体取决于torchmetrics版本和API设计。通常你可能需要自定义逻辑或在所有预测后汇总。
    print(f'Mean Average Precision: {mean_ap_final}')
    return mean_ap_final

if __name__ == '__main__':
    # 获取数据加载器
    dataloader = get_dataloader()

    # 加载 faster_rcnn 模型
    faster_rcnn_model = models.fasterrcnn_resnet50_fpn(weights=models.faster_rcnn.FasterRCNN_ResNet50_FPN_Weights.DEFAULT)
    # 设置模型为评估模式
    faster_rcnn_model.eval()

    # 加载预训练的 SSD300 模型
    ssd_model = models.ssd300_vgg16(weights=models.ssd.SSD300_VGG16_Weights.DEFAULT)
    # 设置模型为评估模式
    ssd_model.eval()

    # models = [yolo_model, yolov10_model, faster_rccn_model, ssd_model]
    # types = ['yolov8', 'yolov10', 'faster_rccn', 'ssd']
    models = [ssd_model]
    types = ['ssd']
    # models = [faster_rcnn_model]
    # types = ['faster_rcnn']
    
    for i in range(len(models)):
        # model_time, model_results = measure_performance(model=models[i], frames=sample_frames, model_type=types[i])
        mean_ap_final = measure_performance(dataloader, models[i], types[i])
        # print(f'{types[i].upper()} Detection Inference Time: {model_time:.2f} seconds per frame')
        # print(f'{types[i].upper()} Detection mAP: {mean_ap_final:.2f} ')
        print(f'{types[i].upper()} Detection mAP: {mean_ap_final} ')


